
import torch
import torch.nn as nn
import timm
class MobileNetV2(nn.Module):
    def __init__(self):
        super(MobileNetV2, self).__init__()
        self.m = timm.create_model('mobilenetv2_100', pretrained=True,features_only=True)



    def forward(self, x):
        x = self.m(x)
        output = x[-3:]

        return tuple(output)

    def init_weights(self):
        pass

if __name__ == '__main__':
    net = MobileNetV2()
    x = torch.randn(2, 3, 224, 224)
    x= net(x)
    for i in x:
        print(i.shape)
